# Written by Bram Cohen
# see LICENSE.txt for license information

from zurllib import urlopen
from urlparse import urlparse
from BT1.btformats import check_message
from BT1.Choker import Choker
from BT1.Storage import Storage
from BT1.StorageWrapper import StorageWrapper
from BT1.FileSelector import FileSelector
from BT1.Uploader import Upload
from BT1.Downloader import Downloader
from BT1.HTTPDownloader import HTTPDownloader
from BT1.Connecter import Connecter
from RateLimiter import RateLimiter
from BT1.Encrypter import Encoder
from BT1.Rerequester import Rerequester
from BT1.DownloaderFeedback import DownloaderFeedback
from RateMeasure import RateMeasure
from CurrentRateMeasure import Measure
from BT1.PiecePicker import PiecePicker
from BT1.Statistics import Statistics
from ConfigDir import ConfigDir
from bencode import bencode, bdecode
from natpunch import UPnP_test
from options import defaults, argslistheader
from BT1.options import defaults as BT1defaults
from os import path, makedirs, listdir
from parseargs import parseargs, formatDefinitions, defaultargs
from threading import Thread, Event
from time import sleep
from clock import clock
from BTcrypto import CRYPTO_OK
from random import shuffle
from BT1.peercache import PeerCache
from utility import exceptionArgsToString
from metadataext import MetadataHandler

from traceback import print_exc

defaults.append(None)
defaults.extend(BT1defaults)

def _failfunc(x):
    print x

class BT1Download:    
    def __init__(self, statusfunc, finfunc, unfinfunc, errorfunc, excfunc, infofunc, doneflag,
                 config, response, infohash, id, rawserver, port,
                 appdataobj = None, private = False, dht = None, ipfilter = None):
        self.statusfunc = statusfunc
        self.finfunc = finfunc
        self.errorfunc = errorfunc
        self.excfunc = excfunc
        self.infofunc = infofunc
        self.doneflag = doneflag
        self.config = config
        self.response = response
        self.infohash = infohash
        self.myid = id
        self.rawserver = rawserver
        self.port = port
        self.ipfilter = ipfilter
        self.private = private
        self.dht = dht
        if dht and not private:
            self.numdhtpeers = 0
        else:
            self.numdhtpeers = -1
        self.dhtpeers = set()
        self.info = self.response['info']
        self.pieces = [self.info['pieces'][x:x+20]
                       for x in xrange(0, len(self.info['pieces']), 20)]
        self.len_pieces = len(self.pieces)
        self.argslistheader = argslistheader
        self.unpauseflag = Event()
        self.unpauseflag.set()
        self.downloader = None
        self.storagewrapper = None
        self.fileselector = None
        self.super_seeding_active = False
        self.filedatflag = Event()
        self.spewflag = Event()
        self.superseedflag = Event()
        self.whenpaused = None
        self.finflag = Event()
        self.rerequest = None
        self.tcp_ack_fudge = config['tcp_ack_fudge']

        self.selector_enabled = config['selector_enabled']
        if appdataobj:
            self.appdataobj = appdataobj
        #elif self.selector_enabled:
        else:
            self.appdataobj = ConfigDir()
            self.appdataobj.deleteOldCacheData( config['expire_cache_data'],
                                                [] )

        self.excflag = self.rawserver.get_exception_flag()
        self.failed = False
        self.checking = False
        self.started = False
        self.peerscache = PeerCache()

        self.picker = PiecePicker(self.len_pieces, config['rarest_first_cutoff'],
                             config['rarest_first_priority_cutoff'])
        self.choker = Choker(config, rawserver.add_task,
                             self.picker, self.finflag.isSet)
        if port is None:
            self.errorfunc("Couldn't listen - no listening port")
            self._failed(None)


    def saveAs(self, filefunc, pathfunc = None):
        try:
            def make(f, forcedir = False):
                if not forcedir:
                    f = path.split(f)[0]
                if f != '' and not path.exists(f):
                    makedirs(f)

            if self.info.has_key('length'):
                file_length = self.info['length']
                file = filefunc(self.info['name'], file_length,
                                self.config['saveas'], False)
                if file is None:
                    return None
                make(file)
                files = [(file, file_length)]
            else:
                file_length = 0L
                for x in self.info['files']:
                    file_length += x['length']
                file = filefunc(self.info['name'], file_length,
                                self.config['saveas'], True)
                if file is None:
                    return None

                # if this path exists, and no files from the info dict exist, we assume it's a new download and 
                # the user wants to create a new directory with the default name
                existing = 0
                if path.exists(file):
                    if not path.isdir(file):
                        self.errorfunc(file + 'is not a dir')
                        return None
                    if len(listdir(file)) > 0:  # if it's not empty
                        for x in self.info['files']:
                            if path.exists(path.join(file, x['path'][0])):
                                existing = 1
                        if not existing:
                            file = path.join(file, self.info['name'])
                            if path.exists(file) and not path.isdir(file):
                                self.errorfunc("Can't create dir - " + self.info['name'])
                                return None
                make(file, True)

                # alert the UI to any possible change in path
                if pathfunc != None:
                    pathfunc(file)

                files = []
                for x in self.info['files']:
                    n = file
                    for i in x['path']:
                        n = path.join(n, i)
                    files.append((n, x['length']))
                    make(n)
        except OSError, e:
            self.errorfunc("Couldn't allocate dir - " + exceptionArgsToString(e))
            return None

        self.filename = file
        self.files = files
        self.datalength = file_length

        return file
    

    def getFilename(self):
        return self.filename

    def _unfinished(self):
        self.finflag.clear()
        self.unfinfunc()

    def _finished(self):
        self.finflag.set()
        if self.storagewrapper.amount_left == 0:
            try:
                self.storage.set_readonly()
            except (IOError, OSError), e:
                self.errorfunc('trouble setting readonly at end - ' + exceptionArgsToString(e))
            if self.superseedflag.isSet():
                self._set_super_seed()
            self.choker.set_round_robin_period(
                max( self.config['round_robin_period'],
                     self.config['round_robin_period'] *
                                         self.info['piece length'] / 200000 ) )
        self.rerequest_complete()
        self.finfunc()

    def _data_flunked(self, amount, index):
        self.ratemeasure_datarejected(amount)
        if not self.doneflag.isSet():
            self.errorfunc('piece %d failed hash check, re-downloading it' % index)

    def _failed(self, reason):
        self.failed = True
        self.doneflag.set()
        if reason is not None:
            self.errorfunc(reason)
        

    def initFiles(self, old_style = False, statusfunc = None):
        if self.doneflag.isSet():
            return None
        if not statusfunc:
            statusfunc = self.statusfunc

        disabled_files = None
        data = self.appdataobj.getTorrentData(self.infohash)

        if self.selector_enabled:
            self.priority = self.config['priority']
            if self.priority:
                try:
                    self.priority = self.priority.split(',')
                    assert len(self.priority) == len(self.files)
                    self.priority = [int(p) for p in self.priority]
                    for p in self.priority:
                        assert p >= -1
                        assert p <= 2
                except:
                    self.errorfunc('bad priority list given, ignored')
                    self.priority = None
            try:
                d = data['resume data']['priority']
                assert len(d) == len(self.files)
                disabled_files = [x == -1 for x in d]
            except:
                try:
                    disabled_files = [x == -1 for x in self.priority]
                except:
                    pass

        try:
            try:
                self.storage = Storage(self.files, self.info['piece length'],
                                       self.doneflag, self.config, disabled_files)
            except IOError, e:
                self.errorfunc('trouble accessing files - ' + exceptionArgsToString(e))
                return None
            if self.doneflag.isSet():
                return None

            self.storagewrapper = StorageWrapper(self.storage, self.config['download_slice_size'],
                self.pieces, self.info['piece length'], self._finished, self._failed,
                statusfunc, self.doneflag, self.config['check_hashes'],
                self._data_flunked, self.rawserver.add_task,
                self.config, self.unpauseflag)
            
        except ValueError, e:
            self._failed('bad data - ' + exceptionArgsToString(e))
        except IOError, e:
            self._failed('IOError - ' + exceptionArgsToString(e))
        if self.doneflag.isSet():
            return None

        if self.selector_enabled:
            self.fileselector = FileSelector(self.files, self.info['piece length'],
                                             self.appdataobj.getPieceDir(self.infohash),
                                             self.storage, self.storagewrapper,
                                             self.rawserver.add_task,
                                             self._failed)
            if data:
                data = data.get('resume data')
                if data:
                    self.fileselector.unpickle(data)
                
        self.checking = True
        if old_style:
            return self.storagewrapper.old_style_init()
        return self.storagewrapper.initialize


    def getCachedTorrentData(self):
        return self.appdataobj.getTorrentData(self.infohash)

    def _make_upload(self, connection, ratelimiter, totalup):
        return Upload(connection, ratelimiter, totalup,
                      self.choker, self.storagewrapper, self.picker,
                      self.config)

    def _kick_peer(self, connection):
        def k(connection = connection):
            connection.close()
        self.rawserver.add_task(k,0)

    def _ban_peer(self, ip):
        self.encoder_ban(ip)
        if self.config['security']:
            self.infofunc(_("Banned IP: ") + ip)

    def _received_raw_data(self, x):
        if self.tcp_ack_fudge:
            x = int(x*self.tcp_ack_fudge)
            self.ratelimiter.adjust_sent(x)

    def _received_data(self, x):
        self.downmeasure.update_rate(x)
        self.ratemeasure.data_came_in(x)

    def _received_http_data(self, x):
        self.downmeasure.update_rate(x)
        self.ratemeasure.data_came_in(x)
        self.downloader.external_data_received(x)

    def _cancelfunc(self, pieces):
        self.downloader.cancel_piece_download(pieces)
        self.httpdownloader.cancel_piece_download(pieces)
   
    def _reqmorefunc(self, pieces):
        self.downloader.requeue_piece_download(pieces)

    def startEngine(self, feedbackfunc = None, finishfunc = None, ratelimiter = None):
        if self.doneflag.isSet():
            return False

        self.checking = False

        if not CRYPTO_OK and self.config['crypto_allowed']:
            self.errorfunc('warning - pycrypto not installed')
            self.config['crypto_allowed'] = 0
        if not self.config['crypto_allowed']:
            self.config['crypto_only'] = 0
        if not self.config['crypto_only']:
            self.config['crypto_stealth'] = 0

        for i in xrange(self.len_pieces):
            if self.storagewrapper.do_I_have(i):
                self.picker.complete(i)
        self.upmeasure = Measure(self.config['max_rate_period'],
                            self.config['upload_rate_fudge'])
        self.downmeasure = Measure(self.config['max_rate_period'])

        if ratelimiter:
            self.ratelimiter = ratelimiter
        else:
            self.ratelimiter = RateLimiter(self.rawserver.add_task,
                                           self.config['upload_unit_size'],
                                           self.config['average_ping'],
                                           self.setConns)
            self.ratelimiter.set_upload_rate(self.config['max_upload_rate'])
        
        self.ratemeasure = RateMeasure()
        self.ratemeasure_datarejected = self.ratemeasure.data_rejected

        self.downloader = Downloader(self.storagewrapper, self.picker,
            self.config['request_backlog'], self.config['max_rate_period'],
            self.len_pieces, self.config['download_slice_size'],
            self._received_data, self.config['snub_time'], self.config['auto_kick'],
            self._kick_peer, self._ban_peer)
        self.downloader.set_download_rate(self.config['max_download_rate'])

        self.connecter = Connecter(self, self.choker, self.upmeasure, self.config, self.ratelimiter,
                                   self._make_upload, self.downloader, self.len_pieces, self.rawserver.add_task)

        self.encoder = Encoder(self.connecter, self.rawserver,
            self.myid, self.config['max_message_length'], self.rawserver.add_task,
            self.config['keepalive_interval'], self.infohash,
            self._received_raw_data, self.config, peerscache = self.peerscache, bans = self.ipfilter)
        self.encoder_ban = self.encoder.ban

        self.httpdownloader = HTTPDownloader(self.storagewrapper, self.picker,
            self.rawserver, self.finflag, self.errorfunc, self.downloader,
            self.config['max_rate_period'], self.infohash, self._received_http_data,
            self.connecter.got_piece)
        if self.response.has_key('httpseeds') and not self.finflag.isSet():
            for u in self.response['httpseeds']:
                self.httpdownloader.make_download(u)

        if self.selector_enabled:
            self.fileselector.tie_in(self.picker, self._cancelfunc,
                    self._reqmorefunc, self.rerequest_ondownloadmore)
            if self.priority:
                self.fileselector.set_priorities_now(self.priority)
            self.appdataobj.deleteTorrentData(self.infohash)
        
        if self.config['super_seeder']:
            self.set_super_seed()

        self.started = True
        return True


    def rerequest_complete(self):
        if self.rerequest:
            self.rerequest.announce(1)

    def rerequest_stopped(self):
        if self.rerequest:
            self.rerequest.announce(2)

    def rerequest_lastfailed(self):
        if self.rerequest:
            return self.rerequest.last_failed
        return False
    
    def rerequest_current(self):
        if self.rerequest:
            return self.rerequest.currentTracker, self.rerequest.currentStatus
        return '', ''

    def rerequest_ondownloadmore(self):
        if self.rerequest:
            self.rerequest.hit()
        # roee88: fileselctor
        self._unfinished()
        
    def startRerequester(self, seededfunc = None, force_rapid_update = False):
        if self.dht and not self.private:
            self.numdhtpeers = 0
        else:
            self.numdhtpeers = -1

        if self.response.has_key('announce-list'):
            trackerlist = self.response['announce-list']
        elif self.response.has_key('announce'):
            trackerlist = [[self.response['announce']]]
        else:
            trackerlist = [] #TODO: this is stupid..

        # Peers Cache
        data = self.getCachedTorrentData()
        if data and data.has_key('peers'):
            self.peerscache.add(data.get('peers'))

        self.rerequest = Rerequester(self.port, self.myid, self.infohash, 
            trackerlist, self.config, 
            self.rawserver.add_task, self.rawserver.add_task,
            self.errorfunc, self.excfunc,
            self.encoder.start_connections,
            self.connecter.how_many_connections, 
            self.storagewrapper.get_amount_left, 
            self.storagewrapper.get_amount_flunked,
            self.upmeasure.get_total, self.downmeasure.get_total,
            self.upmeasure.get_rate, self.downmeasure.get_rate,
            self.doneflag, self.unpauseflag, seededfunc, force_rapid_update,
            self.peerscache )

        self.rerequest.start()

        if self.dht is not None:
            if self.response.has_key('nodes'):
                self.dht.addContacts(self.response['nodes'])                
            self.dhtAnnounce()

    def _init_stats(self):
        self.statistics = Statistics(self.upmeasure, self.downmeasure,
                    self.connecter, self.httpdownloader, self.ratelimiter,
                    self.rerequest_lastfailed, self.filedatflag, self.rerequest_current, lambda : self.numdhtpeers)
        if self.info.has_key('files'):
            self.statistics.set_dirstats(self.files, self.info['piece length'])
        if self.config['spew']:
            self.spewflag.set()

    def autoStats(self, displayfunc = None):
        if not displayfunc:
            displayfunc = self.statusfunc

        self._init_stats()
        DownloaderFeedback(self.choker, self.httpdownloader, self.rawserver.add_task,
            self.upmeasure.get_rate, self.downmeasure.get_rate,
            self.ratemeasure, self.storagewrapper.get_stats,
            self.datalength, self.finflag, self.spewflag, self.statistics,
            displayfunc, self.config['display_interval'])

    def startStats(self):
        self._init_stats()
        d = DownloaderFeedback(self.choker, self.httpdownloader, self.rawserver.add_task,
            self.upmeasure.get_rate, self.downmeasure.get_rate,
            self.ratemeasure, self.storagewrapper.get_stats,
            self.datalength, self.finflag, self.spewflag, self.statistics)
        return d.gather


    def getPortHandler(self):
        return self.encoder


    def shutdown(self, torrentdata = None):
        if torrentdata is None:
            torrentdata = {}
            
        if self.checking or self.started:
            # Wait briefly for files to be released
            sleep(0.2)
            self.storagewrapper.sync()
            self.storage.close()
            self.rerequest_stopped()
            if self.started:
                self.httpdownloader.shutdown()
                self.encoder.shutdown()
                self.connecter.shutdown()
                self.downloader.shutdown()
                self.dhtpeers = set()

        if self.fileselector and (self.started or self.checking):
            if not self.failed:
                self.fileselector.finish()
                torrentdata['resume data'] = self.fileselector.pickle()

        if self.peerscache is not None:
            torrentdata['peers'] = self.peerscache.get(save=True)
            
        try:
            self.appdataobj.writeTorrentData(self.infohash, torrentdata)
        except:
            self.appdataobj.deleteTorrentData(self.infohash) # clear it
            print_exc()                                      # print exception
        self.started = False
        return not self.failed and not self.excflag.isSet()
        # if returns false, you may wish to auto-restart the torrent


    def setUploadRate(self, rate):
        try:
            def s(self = self, rate = rate):
                self.config['max_upload_rate'] = rate
                self.ratelimiter.set_upload_rate(rate)
            self.rawserver.add_task(s)
        except AttributeError:
            pass

    def setConns(self, conns, conns2 = None):
        if not conns2:
            conns2 = conns
        try:
            def s(self = self, conns = conns, conns2 = conns2):
                self.config['min_uploads'] = conns
                self.config['max_uploads'] = conns2
                if (conns > 30):
                    self.config['max_initiate'] = conns + 10
            self.rawserver.add_task(s)
        except AttributeError:
            pass
        
    def setDownloadRate(self, rate):
        try:
            def s(self = self, rate = rate):
                self.config['max_download_rate'] = rate
                self.downloader.set_download_rate(rate)
            self.rawserver.add_task(s)
        except AttributeError:
            pass

    def startConnection(self, ip, port, id):
        self.encoder._start_connection((ip, port), id)
      
    def _startConnection(self, ipandport, id):
        self.encoder._start_connection(ipandport, id)
        
    def setInitiate(self, initiate):
        try:
            def s(self = self, initiate = initiate):
                self.config['max_initiate'] = initiate
            self.rawserver.add_task(s)
        except AttributeError:
            pass

    def getConfig(self):
        return self.config

    def getDefaults(self):
        return defaultargs(defaults)

    def getUsageText(self):
        return self.argslistheader

    def reannounce(self, special = None):
        try:
            def r(self = self, special = special):
                if special is None:
                    self.rerequest.announce()
                else:
                    self.rerequest.announce(specialurl = special)

                self.dhtAnnounce()

            self.rawserver.add_task(r)
        except AttributeError:
            pass

    def dhtAnnounce(self):
        def dhtAnnounceResponse(peers):
            self.dhtpeers.update([peer[0][0] for peer in peers])
            self.numdhtpeers = len(self.dhtpeers)
            self.encoder.start_connections(peers)
        # respect the private torrent flag
        if not self.private and self.dht is not None and self.rerequest.howmany() < self.rerequest.maxpeers:
            self.dht.getPeersAndAnnounce(self.infohash, self.port, dhtAnnounceResponse)
    
    def ChangeTrackerlist(self, trackerlist):
        newtrackerlist = []        
        for tier in trackerlist:
            if len(tier)>1:
                shuffle(tier)
            newtrackerlist += [tier]
        self.rerequest.trackerlist = newtrackerlist
        self.reannounce()
        
    def getResponse(self):
        try:
            return self.response
        except:
            return None

    def Pause(self):
        if not self.storagewrapper:
            return False
        self.unpauseflag.clear()
        self.rawserver.add_task_prio(self.onPause)
        return True

    def onPause(self):
        self.whenpaused = clock()
        if not self.downloader:
            return
        self.downloader.pause(True)
        self.encoder.pause(True)
        self.choker.pause(True)
    
    def Unpause(self):
        self.unpauseflag.set()
        self.rawserver.add_task(self.onUnpause)

    def onUnpause(self):
        if not self.downloader:
            return
        self.downloader.pause(False)
        self.encoder.pause(False)
        self.choker.pause(False)
        #if self.rerequest and self.whenpaused and clock()-self.whenpaused > 60:
        #    self.rerequest.announce(3)      # rerequest automatically if paused for >60 seconds
        if self.rerequest and self.rerequest.stopped:
            # Restarts the rerequester so that the rerequester min interval cycle is re-synchronized on the resume
            self.rerequest.start(3)

    def set_super_seed(self):
        try:
            self.superseedflag.set()
            def s(self = self):
                if self.finflag.isSet():
                    self._set_super_seed()
            self.rawserver.add_task(s)
        except AttributeError:
            pass

    def _set_super_seed(self):
        if not self.super_seeding_active:
            self.super_seeding_active = True
            self.infofunc('** SUPER-SEED OPERATION ACTIVE ** please set Max Uploads so each peer gets 6-8 KB/s')
            def s(self = self):
                self.downloader.set_super_seed()
                self.choker.set_super_seed()
            self.rawserver.add_task(s)
            if self.finflag.isSet():        # mode started when already finished
                def r(self = self):
                    self.rerequest.announce(3)  # so after kicking everyone off, reannounce
                self.rawserver.add_task(r)

    def am_I_finished(self):
        return self.finflag.isSet()

    def get_transfer_stats(self):
        return self.upmeasure.get_total(), self.downmeasure.get_total()
    
def parse_params(params, presets = {}):
    if len(params) == 0:
        return None
    config, args = parseargs(params, defaults, 0, 1, presets = presets)
    if args:
        if config['responsefile'] or config['url']:
            raise ValueError,'must have responsefile or url as arg or parameter, not both'
        if path.isfile(args[0]):
            config['responsefile'] = args[0]
        else:
            try:
                urlparse(args[0])
            except:
                raise ValueError, 'bad filename or url'
            config['url'] = args[0]
    elif (config['responsefile'] == '') == (config['url'] == ''):
        raise ValueError, 'need responsefile or url, must have one, cannot have both'
    return config

# old-style downloader

def get_usage(defaults = defaults, cols = 100, presets = {}):
    return (argslistheader + formatDefinitions(defaults, cols, presets))


def get_response(file, url, errorfunc):
    try:
        if file:
            h = open(file, 'rb')
            try:
                line = h.read(10)   # quick test to see if responsefile contains a dict
                front,garbage = line.split(':',1)
                assert front[0] == 'd'
                int(front[1:])
            except:
                errorfunc(file+' is not a valid responsefile')
                return None
            try:
                h.seek(0)
            except:
                try:
                    h.close()
                except:
                    pass
                h = open(file, 'rb')
        else:
            try:
                h = urlopen(url)
            except:
                errorfunc(url+' bad url')
                return None
        response = h.read()
    
    except IOError, e:
        errorfunc('problem getting response info - ' + str(e))
        return None
    try:    
        h.close()
    except:
        pass
    try:
        try:
            response = bdecode(response)
        except:
            errorfunc("warning: bad data in responsefile")
            response = bdecode(response, sloppy=1)
        check_message(response)
    except ValueError, e:
        errorfunc("got bad file info - " + str(e))
        return None

    return response
